import abc
import datetime
import functools
import logging

from absl import app

import experimental.prophet as prophet
from experimental.prophet.craps.interval import feature


class AlphaBase:
  __metaclass__ = abc.ABCMeta

  @abc.abstractproperty
  def name(self):
    raise NotImeplentedError()

  @abc.abstractmethod
  def get_required(self):
    raise NotImeplentedError()

  @abc.abstractmethod
  def run(self, *mat):
    raise NotImeplentedError()


class KimchiPremium(AlphaBase):
  def __init__(self, target_exchange, ref_exchange):
    self._target_exchange = target_exchange
    self._ref_exchange = ref_exchange

  @property
  def name(self):
    return 'KimchiPremiumAlpha_%s_%s' % (self._target_exchange, self._ref_exchange)

  def get_required(self):
    kimchi_mat_name = 'kimchi_%s_%s' % (self._target_exchange.lower(), self._ref_exchange.lower())
    return [kimchi_mat_name]

  def run(self, premiums):
    row = []

    mean_premium = sum(premiums) / len(premiums)

    sqr_diff_sum = prophet.constant(0.)
    diffs = []
    for premium in premiums:
      diff = premium - mean_premium
      diffs.append(diff)
      sqr_diff_sum += diff * diff

    for premium, diff in zip(premiums, diffs):
      symbol = premium.name.split('/')[1]
      norm_diff = diff / prophet.sqrt(sqr_diff_sum)
      row.append((norm_diff).named_ref(symbol))
    return row


def run_alpha(alpha, interval, train=True):
  ts = prophet.global_timestamp()
  timer = prophet.timer(interval)

  req = []
  requred_features = alpha.get_required()
  for name in requred_features:
    feature_func = feature.FEATURE_INSTANCES[name]
    with prophet.namespace(name):
      row = feature_func(interval)
      req.append(row)

  with prophet.control_if(ts % prophet.to_duration_int(interval) == 0):
    alpha_row = alpha.run(*req)

    if train:
      aggregator = prophet.scalar_value_aggregator([ts, *alpha_row])
      return aggregator
    else:
      return alpah_row


def main(argv):
  alpha_instance = KimchiPremium('Upbit', 'Okex')

  alpha_mat = prophet.run_from_fastfeed_multiprocess(functools.partial(
      run_alpha, alpha_instance, '10s'),
                                                     datetime.date(2019, 8, 1),
                                                     datetime.datetime(2019, 8, 1, 2),
                                                     machine='feed-01.ap-northeast-1.aws')

  print(alpha_instance.name)
  print(alpha_mat)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG')
  app.run(main)
